import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# 设置中文字体支持中文显示（可选）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 模拟数据：教师模型、有老师教的学生模型、无老师教的学生模型的 Loss 和 Acc@1
epochs = list(range(1, 21))

teacher_loss = [
    1.5897, 0.26995, 0.16938, 0.13528, 0.1184,
    0.081307, 0.073402, 0.060262, 0.02126, 0.02895,
    0.014085, 0.013542, 0.012144, 0.013055, 0.013542,
    0.012562, 0.012287, 0.013055, 0.00588, 0.012562
]

teacher_acc = [
    0.5082, 0.9118, 0.9483, 0.9593, 0.9639,
    0.9810, 0.9814, 0.9855, 0.9942, 0.9920,
    0.9953, 0.9957, 0.9967, 0.9972, 0.9967,
    0.9974, 0.9972, 0.9972, 0.9993, 0.9974
]

student_with_teacher_loss = [
    90.075, 0.79687, 0.14372, 0.084808, 0.060262,
    0.047065, 0.040262, 0.031946, 0.02126, 0.02895,
    0.014085, 0.013542, 0.012144, 0.013055, 0.013542,
    0.012562, 0.012287, 0.013055, 0.00588, 0.012562
]

student_with_teacher_acc = [
    0.4022, 0.9596, 0.9723, 0.9799, 0.9855,
    0.9861, 0.9872, 0.9868, 0.9942, 0.9920,
    0.9953, 0.9957, 0.9967, 0.9972, 0.9967,
    0.9974, 0.9972, 0.9972, 0.9993, 0.9974
]

student_without_teacher_loss = [
    101.73, 0.79687, 0.14372, 0.084808, 0.060262,
    0.047065, 0.040262, 0.031946, 0.02126, 0.02895,
    0.014085, 0.013542, 0.012144, 0.013055, 0.013542,
    0.012562, 0.012287, 0.013055, 0.00588, 0.012562
]

student_without_teacher_acc = [
    0.5130, 0.9596, 0.9723, 0.9799, 0.9855,
    0.9861, 0.9872, 0.9868, 0.9942, 0.9920,
    0.9953, 0.9957, 0.9967, 0.9972, 0.9967,
    0.9974, 0.9972, 0.9972, 0.9993, 0.9974
]

# 创建 DataFrame
df_teacher = pd.DataFrame({
    'Epoch': epochs,
    'Model': ['Teacher'] * 20,
    'Loss': teacher_loss,
    'Acc@1': teacher_acc
})

df_student_with_teacher = pd.DataFrame({
    'Epoch': epochs,
    'Model': ['Student with Teacher'] * 20,
    'Loss': student_with_teacher_loss,
    'Acc@1': student_with_teacher_acc
})

df_student_without_teacher = pd.DataFrame({
    'Epoch': epochs,
    'Model': ['Student without Teacher'] * 20,
    'Loss': student_without_teacher_loss,
    'Acc@1': student_without_teacher_acc
})

# 合并数据
df = pd.concat([df_teacher, df_student_with_teacher, df_student_without_teacher], ignore_index=True)

# 保存为 Excel（可选）
df.to_excel("model_comparison.xlsx", index=False)

# 可视化设置
width = 0.25
x_indexes = np.arange(len(epochs))

# ================== Loss 条形图 ==================
plt.figure(figsize=(14, 6))
plt.bar(x_indexes - width, df_teacher['Loss'], width=width, label='Teacher', color='#FF6F61')
plt.bar(x_indexes, df_student_with_teacher['Loss'], width=width, label='Student with Teacher', color='#6B5B95')
plt.bar(x_indexes + width, df_student_without_teacher['Loss'], width=width, label='Student without Teacher', color='#0096FF')

plt.title('Loss Comparison Across Epochs')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.xticks(x_indexes, epochs)
plt.legend()
plt.grid(True, linestyle='--', alpha=0.5)
plt.tight_layout()
plt.savefig("loss_comparison_bar.png")
plt.show()

# ================== Accuracy 条形图 ==================
plt.figure(figsize=(14, 6))
plt.bar(x_indexes - width, df_teacher['Acc@1'], width=width, label='Teacher', color='#FF6F61')
plt.bar(x_indexes, df_student_with_teacher['Acc@1'], width=width, label='Student with Teacher', color='#6B5B95')
plt.bar(x_indexes + width, df_student_without_teacher['Acc@1'], width=width, label='Student without Teacher', color='#0096FF')

plt.title('Accuracy (Acc@1) Comparison Across Epochs')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.xticks(x_indexes, epochs)
plt.legend()
plt.grid(True, linestyle='--', alpha=0.5)
plt.tight_layout()
plt.savefig("accuracy_comparison_bar.png")
plt.show()